"""
build_feature
"""
import numpy as np

import mindspore.nn as nn
import mindspore.numpy as msnp
from mindspore import ops


class BuildFeature(nn.Cell):
    """
    流体力学特征构建模块：基于原始流场数据计算湍流模型所需的物理特征量
    包含涡量、应变率、壁面距离相关参数等流体力学关键特征的计算
    """
    def __init__(self):
        super(BuildFeature, self).__init__()
        # 初始化MindSpore操作符
        self.mul = ops.Mul()          # 乘法操作
        self.div = ops.Div()          # 除法操作
        self.sqrt = ops.Sqrt()        # 开平方操作
        self.add = ops.Add()          # 加法操作
        self.square = ops.Square()    # 平方操作
        self.tanh = ops.Tanh()        # 双曲正切函数
        self.pow = ops.Pow()          # 幂运算
        self.concat = ops.Concat(1)   # 张量拼接（沿维度1）
        self.abs = ops.Abs()          # 绝对值
        self.sign = ops.Sign()        # 符号函数
        self.tan = ops.Tan()          # 正切函数

    def construct(self, mut_value, aoa, reynolds, dis, p_value, ru_value, y_value,
                  u_value, v_value, ux_value, uy_value, vx_value, vy_value):
        """
        构建流体力学特征张量
        
        Args:
            mut_value: 涡粘性系数
            aoa: 攻角(Angle of Attack)，单位为度
            reynolds: 雷诺数
            dis: 到壁面的距离
            p_value: 压力
            ru_value: 密度与速度的乘积（可能为动量）
            y_value: 空间y坐标
            u_value: x方向速度分量
            v_value: y方向速度分量
            ux_value: u对x的偏导数(∂u/∂x)
            uy_value: u对y的偏导数(∂u/∂y)
            vx_value: v对x的偏导数(∂v/∂x)
            vy_value: v对y的偏导数(∂v/∂y)
        
        Returns:
            tuple: 特征张量(前9个特征)和目标值(最后1个特征)
        """
        # --------------------- 流体力学特征计算 ---------------------
        # 计算涡量（速度梯度的旋度大小）
        # 二维涡量公式：ω = ∂v/∂x - ∂u/∂y，此处计算其归一化值
        r_norm = self.sqrt(0.5 * self.square(uy_value - vx_value))

        # 计算与壁面距离相关的特征量
        # 公式包含壁面距离、涡量和双曲正切函数，用于刻画近壁面湍流阻尼效应
        da_r = self.mul(dis, self.mul(dis, self.mul(r_norm, (1 - self.tanh(dis)))))

        # 计算雷诺应力对应的应变率范数
        # 二维应变率张量不变量：√(S_xx² + S_yy² + 0.5*S_xy²)
        s_norm = self.sqrt(self.square(ux_value) + self.square(vy_value) + 0.5 * self.square(uy_value + vx_value))

        # 计算流体熵特征（基于多方气体假设）
        # 公式：1.4*p/ρ^1.4 - 1，其中1.4为空气等熵指数
        entropy = 1.4 * self.div(p_value, self.pow(ru_value, 1.4)) - 1

        # 计算壁面距离相关的指数函数特征
        dref0 = 1.0 / (self.sqrt(reynolds))  # 基于雷诺数的特征长度
        dref1 = ((dis + dref0) - self.abs(dis - dref0)) / 2  # 壁面距离阈值处理
        dref2 = ((dis + dref0) + self.abs(dis - dref0)) / 2
        dsturb_min = dis.min()  # 最小壁面距离
        # 指数函数用于模拟近壁面湍流转捩效应
        expfunc = self.pow(2.71828, self.sqrt(dref1 / dsturb_min))
        expfunc = self.mul(expfunc, self.sqrt(dref0 / dref2)) - 2

        # 计算速度方向特征（弧度制）
        sig = self.sign(y_value)  # y坐标符号，用于确定方向
        # 反正切函数计算速度矢量与x轴夹角
        v_direct = msnp.arctan(self.mul(sig, v_value / u_value))

        # 计算速度在攻角方向的投影
        # 将攻角转换为弧度，并计算速度在攻角方向的分量
        proj_stream = self.mul(sig, self.add(-v_value, self.mul(u_value, self.tan((np.pi * aoa) / 180.0))))

        # 涡粘性系数的尺度修正
        # 基于雷诺数对涡粘性进行归一化，并结合壁面距离的幂函数衰减
        mut_value = mut_value / (reynolds / 1e6)
        trans = 1.0 / self.pow(dis, 0.6)
        mut_value = self.mul(mut_value, trans)

        # --------------------- 特征拼接与输出 ---------------------
        # 将所有特征沿维度1拼接成张量
        df_data = self.concat([msnp.expand_dims(u_value, 1), msnp.expand_dims(r_norm, 1),
                               msnp.expand_dims(s_norm, 1), msnp.expand_dims(entropy, 1),
                               msnp.expand_dims(da_r, 1), msnp.expand_dims(v_direct, 1),
                               msnp.expand_dims(proj_stream, 1), msnp.expand_dims(expfunc, 1),
                               msnp.expand_dims(mut_value, 1)])

        # 返回前9个特征作为输入特征，最后1个特征作为目标值
        return df_data[:, :-1], df_data[:, -1]